#!/bin/bash

################################################################################
# 数据库恢复脚本
# 用途: 从备份恢复 MongoDB, PostgreSQL, Redis 数据
# 作者: AI Writing Platform Team
# 日期: 2025-10-14
################################################################################

set -e  # 遇到错误立即退出

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 配置
BACKUP_ROOT="/mnt/BackupData/aipaper/databases"

# MongoDB 配置
MONGO_USER="admin"
MONGO_PASS="newhand"
MONGO_AUTH_DB="admin"

# PostgreSQL 配置
POSTGRES_USER="admin"
POSTGRES_PASS="newhand"
POSTGRES_DB="aipaper_db"

################################################################################
# 函数定义
################################################################################

log_info() {
    echo -e "${GREEN}[INFO]${NC} $(date '+%Y-%m-%d %H:%M:%S') - $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $(date '+%Y-%m-%d %H:%M:%S') - $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $(date '+%Y-%m-%d %H:%M:%S') - $1"
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

# 显示使用说明
show_usage() {
    cat <<EOF
数据库恢复脚本

用法:
  $0 [选项]

选项:
  -h, --help              显示帮助信息
  -l, --list              列出所有可用备份
  -b, --backup TIMESTAMP  指定要恢复的备份时间戳
  -d, --database DB       指定要恢复的数据库 (mongodb/postgresql/redis/all)
  -f, --force             强制恢复，不提示确认
  --no-backup             恢复前不备份当前数据

示例:
  # 列出所有备份
  $0 --list

  # 恢复最新的备份（所有数据库）
  $0

  # 恢复指定时间的备份
  $0 --backup 20251014_110951

  # 只恢复 MongoDB
  $0 --database mongodb

  # 强制恢复，不提示确认
  $0 --force

  # 恢复前不备份当前数据
  $0 --no-backup
EOF
}

# 列出所有备份
list_backups() {
    log_info "可用的备份列表:"
    echo ""
    echo "时间戳              | 大小      | 压缩文件"
    echo "-------------------+-----------+----------------------------------------"

    if [ -d "$BACKUP_ROOT" ]; then
        find "$BACKUP_ROOT" -maxdepth 1 -type d ! -name "backups" -printf "%f\n" | sort -r | while read dir; do
            if [ -n "$dir" ] && [ -d "$BACKUP_ROOT/$dir" ]; then
                size=$(du -sh "$BACKUP_ROOT/$dir" 2>/dev/null | cut -f1)
                tarfile="$BACKUP_ROOT/${dir}.tar.gz"
                if [ -f "$tarfile" ]; then
                    tarsize=$(du -sh "$tarfile" | cut -f1)
                    echo "$dir | $size     | ✓ ${dir}.tar.gz ($tarsize)"
                else
                    echo "$dir | $size     | ✗ 无压缩文件"
                fi
            fi
        done

        # 显示压缩备份（没有对应目录的）
        find "$BACKUP_ROOT" -name "*.tar.gz" -printf "%f\n" | sort -r | while read tarfile; do
            timestamp="${tarfile%.tar.gz}"
            if [ ! -d "$BACKUP_ROOT/$timestamp" ]; then
                size=$(du -sh "$BACKUP_ROOT/$tarfile" | cut -f1)
                echo "$timestamp | -         | ✓ $tarfile ($size)"
            fi
        done
    else
        log_error "备份目录不存在: $BACKUP_ROOT"
        exit 1
    fi
}

# 选择最新备份
select_latest_backup() {
    local latest=$(find "$BACKUP_ROOT" -maxdepth 1 -type d ! -name "backups" -printf "%f\n" | sort -r | head -1)

    if [ -z "$latest" ]; then
        # 尝试查找压缩备份
        latest=$(find "$BACKUP_ROOT" -name "*.tar.gz" -printf "%f\n" | sort -r | head -1 | sed 's/.tar.gz$//')
    fi

    if [ -z "$latest" ]; then
        log_error "没有找到可用的备份"
        exit 1
    fi

    echo "$latest"
}

# 解压备份
extract_backup() {
    local timestamp=$1
    local backup_dir="$BACKUP_ROOT/$timestamp"
    local tar_file="$BACKUP_ROOT/${timestamp}.tar.gz"

    if [ -d "$backup_dir" ]; then
        log_info "备份目录已存在: $backup_dir"
        return 0
    fi

    if [ -f "$tar_file" ]; then
        log_info "解压备份文件: $tar_file"
        cd "$BACKUP_ROOT"
        tar -xzf "$tar_file"
        log_info "✅ 解压完成"
    else
        log_error "备份文件不存在: $tar_file"
        exit 1
    fi
}

# 备份当前数据
backup_current_data() {
    log_warn "恢复前先备份当前数据..."
    /opt/project/aipaper/scripts/backup-databases.sh
    log_info "✅ 当前数据已备份"
}

# 恢复 MongoDB
restore_mongodb() {
    local backup_dir=$1
    local mongodb_backup="$backup_dir/mongodb"

    if [ ! -d "$mongodb_backup" ]; then
        log_warn "MongoDB 备份不存在，跳过恢复"
        return
    fi

    log_step "恢复 MongoDB 数据库..."

    # 复制备份到容器
    docker cp "$mongodb_backup" mongodb:/tmp/restore_backup

    # 获取所有数据库
    local databases=$(ls "$mongodb_backup")

    for db in $databases; do
        if [ -d "$mongodb_backup/$db" ]; then
            log_info "  恢复数据库: $db"

            # 显示集合数量
            local collections=$(ls "$mongodb_backup/$db"/*.bson 2>/dev/null | wc -l)
            log_info "    集合数量: $collections"

            # 执行恢复
            docker exec mongodb mongorestore \
                --username "$MONGO_USER" \
                --password "$MONGO_PASS" \
                --authenticationDatabase "$MONGO_AUTH_DB" \
                --db "$db" \
                --drop \
                "/tmp/restore_backup/$db" \
                2>&1 | grep -E "(done|restored|dropping)" || true

            # 验证恢复结果
            local count=$(docker exec mongodb mongosh \
                --quiet \
                --username "$MONGO_USER" \
                --password "$MONGO_PASS" \
                --authenticationDatabase "$MONGO_AUTH_DB" \
                "$db" \
                --eval "db.stats().collections" 2>/dev/null || echo "0")

            log_info "    ✅ $db 恢复完成 (集合数: $count)"
        fi
    done

    # 清理容器内的临时文件
    docker exec mongodb rm -rf /tmp/restore_backup

    log_info "✅ MongoDB 恢复完成"
}

# 恢复 PostgreSQL
restore_postgresql() {
    local backup_dir=$1
    local pg_backup="$backup_dir/postgresql"

    if [ ! -d "$pg_backup" ]; then
        log_warn "PostgreSQL 备份不存在，跳过恢复"
        return
    fi

    log_step "恢复 PostgreSQL 数据库..."

    # 查找 SQL 文件
    local sql_file="$pg_backup/${POSTGRES_DB}.sql"
    local dump_file="$pg_backup/postgres_backup_*.dump"

    if [ -f "$sql_file" ]; then
        log_info "  使用 SQL 文件恢复: $sql_file"

        # 复制到容器
        docker cp "$sql_file" postgres:/tmp/restore.sql

        # 执行恢复
        docker exec postgres psql \
            -U "$POSTGRES_USER" \
            -d "$POSTGRES_DB" \
            -f /tmp/restore.sql \
            2>&1 | grep -v "ERROR.*already exists" || true

        # 清理
        docker exec postgres rm -f /tmp/restore.sql

        log_info "  ✅ PostgreSQL 恢复完成"
    elif ls $dump_file 1> /dev/null 2>&1; then
        log_info "  使用 DUMP 文件恢复"

        local dump=$(ls $dump_file | head -1)
        docker cp "$dump" postgres:/tmp/restore.dump

        docker exec postgres pg_restore \
            -U "$POSTGRES_USER" \
            -d "$POSTGRES_DB" \
            --clean \
            --if-exists \
            /tmp/restore.dump \
            2>&1 | grep -v "ERROR.*does not exist" || true

        docker exec postgres rm -f /tmp/restore.dump

        log_info "  ✅ PostgreSQL 恢复完成"
    else
        log_warn "未找到 PostgreSQL 备份文件"
    fi
}

# 恢复 Redis
restore_redis() {
    local backup_dir=$1
    local redis_backup="$backup_dir/redis"

    if [ ! -d "$redis_backup" ]; then
        log_warn "Redis 备份不存在，跳过恢复"
        return
    fi

    log_step "恢复 Redis 数据库..."

    local rdb_file="$redis_backup/dump.rdb"

    if [ -f "$rdb_file" ]; then
        log_info "  停止 Redis 服务..."
        docker stop redis > /dev/null

        log_info "  复制 RDB 文件..."
        docker cp "$rdb_file" redis:/data/dump.rdb

        log_info "  启动 Redis 服务..."
        docker start redis > /dev/null

        # 等待 Redis 启动
        sleep 2

        # 验证
        local keys=$(docker exec redis redis-cli DBSIZE | tr -d '\r')
        log_info "  ✅ Redis 恢复完成 (键数量: $keys)"
    else
        log_warn "未找到 Redis RDB 文件"
    fi
}

# 确认操作
confirm_restore() {
    local timestamp=$1
    local database=$2

    echo ""
    log_warn "=========================================="
    log_warn "警告: 此操作将覆盖当前数据库数据"
    log_warn "=========================================="
    echo ""
    echo "备份时间: $timestamp"
    echo "恢复目标: $database"
    echo ""

    read -p "确认继续? (yes/no): " confirm

    if [ "$confirm" != "yes" ]; then
        log_info "操作已取消"
        exit 0
    fi
}

################################################################################
# 主函数
################################################################################

main() {
    local backup_timestamp=""
    local database="all"
    local force=false
    local do_backup=true

    # 解析参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_usage
                exit 0
                ;;
            -l|--list)
                list_backups
                exit 0
                ;;
            -b|--backup)
                backup_timestamp="$2"
                shift 2
                ;;
            -d|--database)
                database="$2"
                shift 2
                ;;
            -f|--force)
                force=true
                shift
                ;;
            --no-backup)
                do_backup=false
                shift
                ;;
            *)
                log_error "未知参数: $1"
                show_usage
                exit 1
                ;;
        esac
    done

    log_info "=========================================="
    log_info "数据库恢复工具"
    log_info "=========================================="

    # 如果未指定备份，使用最新的
    if [ -z "$backup_timestamp" ]; then
        backup_timestamp=$(select_latest_backup)
        log_info "使用最新备份: $backup_timestamp"
    fi

    local backup_dir="$BACKUP_ROOT/$backup_timestamp"

    # 解压备份
    extract_backup "$backup_timestamp"

    # 验证备份目录
    if [ ! -d "$backup_dir" ]; then
        log_error "备份目录不存在: $backup_dir"
        exit 1
    fi

    # 确认操作
    if [ "$force" = false ]; then
        confirm_restore "$backup_timestamp" "$database"
    fi

    # 备份当前数据
    if [ "$do_backup" = true ]; then
        backup_current_data
    fi

    # 执行恢复
    case $database in
        mongodb)
            restore_mongodb "$backup_dir"
            ;;
        postgresql)
            restore_postgresql "$backup_dir"
            ;;
        redis)
            restore_redis "$backup_dir"
            ;;
        all)
            restore_mongodb "$backup_dir"
            restore_postgresql "$backup_dir"
            restore_redis "$backup_dir"
            ;;
        *)
            log_error "未知数据库类型: $database"
            exit 1
            ;;
    esac

    log_info "=========================================="
    log_info "✅ 数据库恢复完成"
    log_info "恢复来源: $backup_dir"
    log_info "=========================================="
}

# 运行主函数
main "$@"
